from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import yaml
import os
import time
from captcha_recognition.tesseract_recognizer import TesseractRecognizer
# from captcha_recognition.tensorflow_recognizer import TensorflowRecognizer

class AutoLogin:
    def __init__(self, config_path: str):
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)

        self.driver = webdriver.Chrome()
        self.recognizer = self._init_recognizer()

    def _init_recognizer(self):
        method = self.config['login']['captcha_method']
        if method == 'tesseract':
            return TesseractRecognizer()
        elif method == 'tensorflow':
            return None
            # return TensorflowRecognizer()
        else:
            raise ValueError(f"不支持的验证码识别方法: {method}")

    async def login(self):
        try:
            # 打开登录页面
            self.driver.get(self.config['login']['url'])

            # 等待页面元素加载
            wait = WebDriverWait(self.driver, 10)

            # 定位并输入用户名
            username_input = wait.until(
                EC.presence_of_element_located((By.NAME, "username"))
            )
            username_input.send_keys(self.config['login']['username'])

            # 定位并输入密码
            password_input = self.driver.find_element(By.NAME, "password")
            password_input.send_keys(self.config['login']['password'])

            # 获取验证码图片
            captcha_img = self.driver.find_element(By.CLASS_NAME, "captcha-image")
            captcha_img.screenshot("temp_captcha.png")

            # 识别验证码
            captcha_text = await self.recognizer.recognize("temp_captcha.png")

            # 输入验证码
            captcha_input = self.driver.find_element(By.NAME, "captcha")
            captcha_input.send_keys(captcha_text)

            # 点击登录按钮
            login_button = self.driver.find_element(By.CLASS_NAME, "login-button")
            login_button.click()

            # 等待登录成功
            time.sleep(2)

            return True

        except Exception as e:
            print(f"登录失败: {str(e)}")
            return False

        finally:
            # 清理临时文件
            if os.path.exists("temp_captcha.png"):
                os.remove("temp_captcha.png")

    def close(self):
        self.driver.quit()
